package main

import (
	"encoding/gob"
	"flag"
	"fmt"
	"log"
	"os"
	"os/signal"
	"strconv"
	"syscall"
	"time"

	"honeybadgerbft/hbbft"
)

func main() {
	// 命令行参数
	var (
		nodeID     = flag.Int("id", 0, "Node ID")
		selfAddr   = flag.String("addr", "localhost:8000", "Self address")
		peers      = flag.String("peers", "1:localhost:8001,2:localhost:8002,3:localhost:8003", "Peer addresses (id:addr,id:addr)")
		nodes      = flag.Int("n", 4, "Total number of nodes")
		faulty     = flag.Int("f", 1, "Number of faulty nodes")
		batchSize  = flag.Int("batch", 100, "Batch size")
	)
	flag.Parse()

	// 注册gob类型
	registerGobTypes()

	// 解析peer地址
	peerMap, err := parsePeers(*peers)
	if err != nil {
		log.Fatalf("Failed to parse peers: %v", err)
	}

	// 创建配置
	config := hbbft.Config{
		N:              *nodes,
		F:              *faulty,
		ID:             uint64(*nodeID),
		Nodes:          make([]uint64, *nodes),
		BatchSize:      *batchSize,
		SelfAddr:       *selfAddr,
		Peers:          peerMap,
		SkipProofVerify: true, // 论文实验简化模式（跳过默克尔证明、阈值加密等）
	}

	// 填充节点ID列表
	for i := 0; i < *nodes; i++ {
		config.Nodes[i] = uint64(i)
	}

	log.Printf("Starting node %d on %s", *nodeID, *selfAddr)
	log.Printf("Peers: %v", peerMap)

	// 创建HoneyBadger实例
	hb := hbbft.NewHoneyBadger(config, nil) // 先传nil，稍后设置网络

	// 创建网络
	network := hbbft.CreateTCPNetwork(hbbft.NetworkConfig{
		SelfID:   uint64(*nodeID),
		SelfAddr: *selfAddr,
		Peers:    peerMap,
	}, hb)

	// 设置网络
	hb.SetNetwork(network)

	// 启动网络
	if err := network.Start(); err != nil {
		log.Fatalf("Failed to start network: %v", err)
	}

	// 启动HoneyBadger
	if err := hb.Start(); err != nil {
		log.Fatalf("Failed to start HoneyBadger: %v", err)
	}

	// 添加一些测试交易
	go func() {
		// 初始快速填充交易池（足够多的交易）
		initialTxs := *batchSize * 10 // 增加到10倍batchSize
		log.Printf("Node %d: Initializing transaction pool with %d transactions", *nodeID, initialTxs)
		for i := 0; i < initialTxs; i++ {
			tx := hbbft.SimpleTransaction(fmt.Sprintf("test-tx-%d-%d", *nodeID, i))
			hb.AddTransaction(tx)
		}
		
		// 持续添加交易保持池不为空
		ticker := time.NewTicker(50 * time.Millisecond)
		defer ticker.Stop()
		txCounter := initialTxs
		for {
			select {
			case <-ticker.C:
				// 每50ms添加batchSize笔交易，确保交易池始终充足
				for j := 0; j < *batchSize; j++ {
					tx := hbbft.SimpleTransaction(fmt.Sprintf("test-tx-%d-%d", *nodeID, txCounter))
					hb.AddTransaction(tx)
					txCounter++
				}
			}
		}
	}()

	// 监控输出
	go func() {
		ticker := time.NewTicker(5 * time.Second)
		defer ticker.Stop()
		for {
			select {
			case <-ticker.C:
				outputs := hb.Outputs()
				if len(outputs) > 0 {
					log.Printf("Node %d committed %d epochs", *nodeID, len(outputs))
					for epoch, txs := range outputs {
						log.Printf("  Epoch %d: %d transactions", epoch, len(txs))
					}
				}
				
                // 输出性能统计（修正延迟格式）
                stats := hb.GetPerformanceStats()
                avgLatencyMs := float64(0)
                if latency, ok := stats["avg_epoch_latency"].(int64); ok {
                    avgLatencyMs = float64(latency)
                }
                log.Printf("Performance - TPS: %.2f, Epochs: %v, Avg Latency: %.2fms", 
                    stats["throughput_tps"], stats["total_epochs"], avgLatencyMs)
			}
		}
	}()

	// 等待信号
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
	<-sigChan

	log.Printf("Shutting down node %d", *nodeID)
	
	// 关闭HoneyBadger并输出最终统计
	hb.Close()
	
	// 关闭网络
	network.Close()
}

// registerGobTypes 注册所有需要序列化的类型
func registerGobTypes() {
	gob.Register(&hbbft.ACSMessage{})
	gob.Register(&hbbft.RBCMessage{})
	gob.Register(&hbbft.BBAMessage{})
	gob.Register(&hbbft.HBMessage{})
	gob.Register(&hbbft.BroadcastMessage{})
	gob.Register(&hbbft.AgreementMessage{})
	gob.Register(&hbbft.ProofRequest{})
	gob.Register(&hbbft.EchoRequest{})
	gob.Register(&hbbft.ReadyRequest{})
	gob.Register(&hbbft.BinaryValueRequest{})
	gob.Register(&hbbft.AuxRequest{})
	gob.Register(hbbft.SimpleTransaction(""))
	// 加密相关类型
	gob.Register(&hbbft.DecryptShareMsg{})
	gob.Register(&hbbft.EncryptedBatch{})
}

// parsePeers 解析peer地址字符串
func parsePeers(peersStr string) (map[uint64]string, error) {
	peers := make(map[uint64]string)
	if peersStr == "" {
		return peers, nil
	}

	// 解析格式: "1:localhost:8001,2:localhost:8002"
	for _, peer := range split(peersStr, ",") {
		parts := split(peer, ":")
		if len(parts) < 2 {
			return nil, fmt.Errorf("invalid peer format: %s", peer)
		}

		id, err := strconv.ParseUint(parts[0], 10, 64)
		if err != nil {
			return nil, fmt.Errorf("invalid peer ID: %s", parts[0])
		}

		// 重新组合地址部分
		addr := parts[1]
		for i := 2; i < len(parts); i++ {
			addr += ":" + parts[i]
		}

		peers[id] = addr
	}

	return peers, nil
}

// split 简单的字符串分割函数
func split(s, sep string) []string {
	var result []string
	start := 0
	for i := 0; i < len(s); i++ {
		if i+len(sep) <= len(s) && s[i:i+len(sep)] == sep {
			result = append(result, s[start:i])
			start = i + len(sep)
			i += len(sep) - 1
		}
	}
	result = append(result, s[start:])
	return result
}
